#!/usr/bin/env python
# coding: utf-8

"""
train.py
create by cuiyu 2025-02-23
说明：
 * ultralytics 模型通用训练文件


用法:
  1. 修改文件中的 G_DEUBG 和相关参数
  2. 运行: train.py
    * 训练日志保存在: log/train.log
    * 训练结果保存在: project/name 目录下
  
更新:
2025-02-23:
"""


import datetime
import sys
import shutil
from pathlib import Path
from loguru import logger as log
from dataclasses import dataclass
from ultralytics import YOLO
from matplotlib import pyplot as plt
from ruamel.yaml import YAML
from app_config import CModelConfig

# 指定使用 SimHei 字体（黑体），支持中文显示
plt.rcParams["font.sans-serif"] = ["simsun"] 
plt.rcParams["axes.unicode_minus"] = False  # 用来正常显示负号


# 全局变量
_APP_FILE = Path(sys.argv[0])
_APP_ROOT = _APP_FILE.parent

# import os
# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:1024"



def modelTrain(train_cfg: CModelConfig, train_file: Path, model_out: Path) -> str:
    """
    模型训练，并将训练结果保存到 model_out 目录下

    Args:
      :train_cfg: CTrainConfig;  训练配置
      :model_out: Path;          模型输出目录

    Returns:
      :str; 训练结果目录 = proj_name
    """

    train_id = f"{datetime.datetime.today():%Y%m%d-%H%M%S}"
    project = "train-out"
    proj_name = "jgjc-" + train_id

    # 创建 YAML 对象
    yaml = YAML(typ="safe")
    yaml.preserve_quotes = True  # 保留字符串引号（可选）
    yaml.indent(mapping=2, sequence=4, offset=2)  # 设置缩进格式
    yaml.allow_unicode = True  # 关键：禁用 Unicode 转义
    with open(train_file, "r", encoding="utf-8") as file:
        # 使用safe_load()来解析YAML文件内容到Python字典
        train_data = yaml.load(file)

    # log start
    log.info(" ")
    log.info(f"训练开始 ID = {train_id}")
    log.info(f"模型文件 = {train_cfg.model_file}")
    log.info(f"训练文件 = {train_file}")
    log.info(f"图片目录 = {train_data['path']}")
    log.info(f"类别 = {train_data['names']}")

    # Load a model
    model = YOLO(train_cfg.model_file)

    # train
    # imgsz=[2048], rect=Ture
    # model.train(data=train_file, imgsz=640, batch=2, epochs=2)

    # 大尺寸适配：2048px长边在保持细节的同时控制显存消耗（4090x2000→2048 + rect=True）
    # 缺陷特性优化：增强曝光扰动突出低对比裂缝，限制旋转角度避免方向失真
    # 计算效率：rect模式减少无效像素填充，AMP加速训练
    # 模型选择：yolo11l-seg在COCO分割任务中mAP50达54.2%，优于m/s型号
    # 建议训练初期使用10%数据验证参数有效性，再扩展至全量数据。若出现显存溢出，可依次尝试：降低batch→关闭mosaic→减小imgsz→启用梯度累积。
    model.train(
        # 必要参数
        data = train_file,

        # 训练终止控制参数
        epochs = train_cfg.epochs,        # 300
        patience = train_cfg.patience,    # 早停机制监测mAP50连续50轮无提升

        # 训练尺度参数
        batch = train_cfg.batch,          # 8
        imgsz = train_cfg.imgsz,          # 控制显存消耗，提升训练速度
        rect = train_cfg.rect,            # 启动矩阵训练，对于 4090✕2000的图像，可以两种拼接为一张

        # 训练保存参数
        project=project,
        name = proj_name,
        plots = True,           # 生成PR曲线和混淆矩阵

        # 训练优化参数
        freeze = train_cfg.freeze,        # 冻结前10层权重，用于小数据集微调
        optimizer = train_cfg.optimizer,  # 大数据集首选 SGD，相较于SGD，AdamW在复杂场景下收敛更稳定，适合裂缝类不规则目标
        lr0 = train_cfg.lr0,              # 初始学习率（实测道路缺陷最佳值）; 初始学习率设为0.001（低于默认0.01），避免大尺寸图像训练初期梯度爆炸
        lrf = train_cfg.lrf,              # 最终学习率（实测道路缺陷最佳值）; 最终学习率设为0.01（低于默认0.1），避免训练后期学习率过快
        momentum = train_cfg.momentum,    # 动量参数防止局部最优
        cos_lr = True,                    # 余弦退火调度
        amp = True,                       # 混合精度训练节省30%显存
        profile = train_cfg.profile,      # 训练期间对 ONNX 和 TensorRT 的速度进行分析，有助于优化模型部署。        

        # 数据增强参数
        mosaic = 0.5,           # mosaic 增强概率, 大尺寸图像可能需降低 mosaic 概率，比如 0.5，避免计算负担。
        close_mosaic = 15,      # 最后15个 epoch 关闭马赛克增强
        degrees = 15.0,         # 旋转角度±15°增强裂缝方向鲁棒性
        shear = 5.0,            # 剪切变换模拟路面扭曲
        hsv_h = 0.1,            # 色调偏移±0.1适应不同光照条件[5](@ref)
        hsv_s = 0.9,            # 饱和度扰动系数        
        hsv_v = 0.7,            # 曝光增强系数（突出低对比度裂缝）[8](@ref)
        overlap_mask = True,    # 允许掩膜重叠（龟裂区域常存在交叉）  

        # 硬件控制参数
        workers = train_cfg.workers,      # 16 数据加载线程数

        # [官方文档中没有，但是可以通过] 验证参数
        # val = True,             # 每 epoch 验证
        # split = 'test',         # 使用独立测试集验证
    )

    # log end
    log.info(f"训练结束: [{proj_name}]")

    # 将训练结果移动到全局的 model_out 目录
    train_out_dir = _APP_ROOT / project / proj_name
    model_out_dir = model_out / proj_name
    if (model_out_dir.exists()):
        shutil.rmtree(model_out_dir)
    shutil.move(train_out_dir, model_out)

    return proj_name

#==== 本文件的测试代码

def _logInit(baseDir: Path) -> None:
    log.add(baseDir / "train.log", rotation="500 MB", encoding="utf-8")

def _main():
    from runtimer import runtimer_create
    from app_config import CAppConfig
    MODEL_ROOT = Path("/aaby/netdisk/mydata/ai-data/model/jgjc")  # 最终模型输出的目录

    # 测试训练
    wc = CAppConfig()
    wc.load(_APP_ROOT / "app-debug.yaml")
    train_file = Path(wc.pipe_cfg.img_root_train).parent / "train.yaml"
    # train_cfg = CTrainConfig()
    # train_cfg.model_file = Path("/aaby/netdisk/mydata/ai-data/model/ultralytics/seg/yolo11n-seg.pt")
    # train_cfg.train_file = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/rddhw/test/train.yaml")
    # train_cfg.batch = 2
    # train_cfg.epochs = 2
    # train_cfg.patience = 0
    # train_cfg.workers = 0
    # train_cfg.freeze = 10   
    # train_cfg.imgsz = 640
    # train_cfg.rect = True
    # train_cfg.profile = True  # 训练期间对 ONNX 和 TensorRT 的速度进行分析，有助于优化模型部署。

    with runtimer_create("训练"):
        modelTrain(wc.model_cfg, train_file, MODEL_ROOT)
    pass


if __name__ == "__main__":
    _logInit(_APP_ROOT / "log/")
    _main()
